﻿using AdvVerifyGUI.Display;
using AdvVerifyGUI.Stage;
using AdvVerifyGUI.UseCases;
using AdvVerifyGUI.Utils.Image;
using AdvVerifyModels.Entities;
using HWModule;
using System;
using System.Diagnostics;
using System.Threading.Tasks;
using System.Windows;

namespace AdvVerifyGUI.Job;

enum PatternKind
{
    High,
    Align,
    Mp,
    Addr,
    PreDose,
    Focus,
}

class ImagingResult
{
    public int Index { get; set; }

    public PatternKind Kind { get; set; }

    public ImageData Image { get; set; }

    public double? AddrMatchDx { get; set; }

    public double? AddrMatchDy { get; set; }

    public double MatchDx { get; set; }

    public double MatchDy { get; set; }

    public int? AddressingScore { get; set; }

    public int? Score { get; set; }

    public Coordinate Coordinate { get; set; }

    public Pattern Pattern { get; set; }
}

class ImagingPipeline(IJobContext context,
    IWaferNavigator nav,
    ISEMClient client,
    IImagingService imagingService)
{
    Coordinate last;

    int lastAddressingId;
    double? addrMatchDx;
    double? addrMatchDy;
    int? addressingScore;
    Vector offset;

    int lastFocusId;

    int index = 0;

    public IWaferNavigator Navigator
    {
        get
        {
            return nav;
        }
        set
        {
            nav = value;
        }
    }

    public MatchFailureMode MatchFailureMode { get; set; }

    public bool IsDYNEnabled { get; set; } = true;

    public bool SkipAutoFocus { get; set; }

    public async Task<ImagingResult> MoveTo(Coordinate c, Pattern pattern)
    {
        if (last != c)
        {
            if (!pattern.UseDifferentialDrive && pattern.Addressing.IsRegistered)
            {
                lastAddressingId = pattern.Id;
                offset = await Address(c, pattern).ConfigureAwait(false);
            }
            else if (pattern.ParentId != lastAddressingId || c.X != last.X || c.Y != last.Y)
            {
                lastAddressingId = 0;
                addrMatchDx = null;
                addrMatchDy = null;
                addressingScore = null;
                offset = default;
            }
        }

        if (pattern.Kind != 0)
        {
            await client.Apply(pattern.Reference.ScanSettings).ConfigureAwait(false);
        }

        if (!SkipAutoFocus)
        {
            var param = pattern.Reference.AutoFocusParam;
            var mode = pattern.Kind == 0 || pattern.UseDifferentialDrive ?
                AutoFocusMode.Disabled : param.Mode;
            switch (mode)
            {
                case AutoFocusMode.Disabled:
                    if (last != c)
                    {
                        await nav.MoveTo(c.X, c.Y, c.Dx, c.Dy, offset.X, offset.Y).ConfigureAwait(false);
                    }
                    break;
                case AutoFocusMode.Once:
                    if (pattern.Id == lastFocusId)
                    {
                        goto case AutoFocusMode.Disabled;
                    }
                    await AutoFocus(c,
                        offset,
                        pattern.Reference.ScanSettings,
                        pattern.Reference.AutoFocusParam).ConfigureAwait(false);
                    lastFocusId = pattern.Id;
                    break;
                case AutoFocusMode.Always:
                    await AutoFocus(c,
                        offset,
                        pattern.Reference.ScanSettings,
                        pattern.Reference.AutoFocusParam).ConfigureAwait(false);
                    lastFocusId = 0;
                    break;
            }
        }

        await Task.Delay(pattern.Reference.ScanSettings.WaitTime).ConfigureAwait(false);

        await PreDose(pattern.Reference.PreDoseOptions).ConfigureAwait(false);

        byte[] b = [];
        ImageRead.GetBitmap(pattern.Reference.ImagePath, ref b);

        ImageData img;
        if (pattern.Kind == 0)
        {
            img = await imagingService.Current.Take().ConfigureAwait(false);
        }
        else if (IsDYNEnabled)
        {
            img = await client.DYNTake().ConfigureAwait(false);
        }
        else
        {
            img = await client.Take().ConfigureAwait(false);
        }
        var (dx, dy, score) = Match(pattern.Reference.PatternRecognition, b, img);
        if (score > pattern.Reference.Threshold)
        {
            if (pattern.Reference.DoCentering && last != c)
            {
                if (Math.Abs(dx / img.Width) > 0.05 || Math.Abs(dy / img.Height) > 0.05)
                {
                    var cond = client.Condition;
                    var mat = Transformation.NewTransformation(cond.PixelSize, cond.ImageAngle);
                    var off = mat.Transform(new Vector(-dx, -dy));
                    await nav.MotionControl.RelativeMove(off.X, off.Y).ConfigureAwait(false);
                    img = await client.Take().ConfigureAwait(false);
                    (dx, dy, score) = Match(pattern.Reference.PatternRecognition, b, img);
                }
            }
        }

        var mp = new ImagingResult
        {
            Index = index,
            Kind = PatternKind.High + pattern.Kind,
            Image = img,
            AddrMatchDx = addrMatchDx,
            AddrMatchDy = addrMatchDy,
            MatchDx = dx,
            MatchDy = dy,
            AddressingScore = addressingScore,
            Score = score,
            Pattern = pattern,
            Coordinate = c,
        };
        await context.Write(mp).ConfigureAwait(false);

        if (score < pattern.Reference.Threshold)
        {
            switch (MatchFailureMode)
            {
                case MatchFailureMode.Terminate:
                    throw new ThresholdException(score, pattern.Reference.Threshold);
            }
        }

        last = c;
        index += 1;
        return mp;
    }

    async Task<Vector> Address(Coordinate c, Pattern pattern)
    {
        await client.Apply(pattern.Addressing.ScanSettings).ConfigureAwait(false);
        if (pattern.Addressing.AutoFocusParam.Mode == AutoFocusMode.Always)
        {
            await AutoFocus(c,
                default,
                pattern.Addressing.ScanSettings,
                pattern.Addressing.AutoFocusParam).ConfigureAwait(false);
            lastFocusId = 0;
        }

        await nav.MoveTo(c.X, c.Y, c.Dx + c.AddrOffsetX, c.Dy + c.AddrOffsetY).ConfigureAwait(false);
        await Task.Delay(pattern.Addressing.ScanSettings.WaitTime).ConfigureAwait(false);

        await PreDose(pattern.Addressing.PreDoseOptions).ConfigureAwait(false);

        byte[] b = [];
        ImageRead.GetBitmap(pattern.Addressing.ImagePath, ref b);

        var cond = client.Condition;
        var img = await client.Take().ConfigureAwait(false);
        var (dx, dy, score) = Match(pattern.Addressing.PatternRecognition, b, img);
        await context.Write(new ImagingResult
        {
            Index = index,
            Kind = PatternKind.Addr,
            Image = img,
            Score = score,
            Pattern = pattern,
            MatchDx = dx,
            MatchDy = dy,
            Coordinate = c,
        }).ConfigureAwait(false);
        if (score < pattern.Addressing.Threshold)
        {
            switch (MatchFailureMode)
            {
                case MatchFailureMode.Terminate:
                    throw new ThresholdException(score, pattern.Addressing.Threshold);
            }
        }

        addrMatchDx = dx;
        addrMatchDy = dy;
        addressingScore = score;
        var mat = Transformation.NewTransformation(cond.PixelSize, cond.ImageAngle);
        return mat.Transform(new Vector(-dx, -dy));
    }

    async Task AutoFocus(Coordinate c, Vector offset, ScanSettings scan, AutoFocusParam param)
    {
        await nav.MoveTo(c.X, c.Y, c.Dx + param.Dx, c.Dy + param.Dy, offset.X, offset.Y).ConfigureAwait(false);
        await Task.Delay(param.WaitTime).ConfigureAwait(false);
        if (param.DotAvg != scan.DotAvg || param.LineAvg != scan.LineAvg || param.FrameAvg != scan.FrameAvg)
            await client.ApplyAvg(param.DotAvg, param.LineAvg, param.FrameAvg).ConfigureAwait(false);
        var cond = client.Condition;
        double[] gunFocus = [cond.Gunfocus_coarsetuning_in, cond.Gunfocus_finecontrol_in];
        var focus = new FocusTrackingUseCase(client);
        await foreach (var (img, quality, current) in focus.FocusSearch(gunFocus).ConfigureAwait(false))
        {
            var result = new ImagingResult
            {
                Kind = PatternKind.Focus,
                Image = img,
            };
            var imgQuality = new ImageQuality(quality, current);
            await context.Write(result, imgQuality).ConfigureAwait(false);
        }

        if (param.DotAvg != scan.DotAvg || param.LineAvg != scan.LineAvg || param.FrameAvg != scan.FrameAvg)
            await client.ApplyAvg(scan.DotAvg, scan.LineAvg, scan.FrameAvg).ConfigureAwait(false);

        await nav.MoveTo(c.X, c.Y, c.Dx, c.Dy, offset.X, offset.Y).ConfigureAwait(false);
    }

    async Task PreDose(PreDoseOptions options)
    {
        if (!options.IsEnabled || options.WaitTime == 0)
        {
            return;
        }

        var cond = client.Condition;
        var settings = client.Settings;
        await settings.SetFrameAverage(5).ConfigureAwait(false);
        await settings.SetFov(options.FOV).ConfigureAwait(false);
        var sw = Stopwatch.StartNew();
        while (sw.ElapsedMilliseconds < options.WaitTime)
        {
            var img = await client.Take().ConfigureAwait(false);
            await context.Write(new ImagingResult
            {
                Index = index,
                Kind = PatternKind.PreDose,
                Image = img,
            }).ConfigureAwait(false);
        }
        await settings.SetFrameAverage(cond.FrameAverage).ConfigureAwait(false);
        await settings.SetFov(cond.Fov).ConfigureAwait(false);
    }

    static (double, double, int) Match(PatternRecognition recognition, byte[] pattern, ImageData img)
    {
        var (w, h, data, _) = img;
        return recognition.Algorithm switch
        {
            PatternMatchAlgorithm.FastMatchNormal => Alignment.Align(w, h, pattern, data, 0),
            PatternMatchAlgorithm.EnhancedMatch => Alignment.Align(w, h, pattern, data, 3),
            PatternMatchAlgorithm.FastROIMatch => ROIMatch(recognition, w, h, pattern, data),
            _ => throw new ArgumentException($"Unknown algorithm {recognition.Algorithm}"),
        };
    }

    static (double, double, int) ROIMatch(PatternRecognition recognition, int w, int h, byte[] pattern, byte[] img)
    {
        var roiX = (int)(recognition.RoiX * w);
        var roiY = (int)(recognition.RoiY * h);
        var roiW = (int)(recognition.RoiW * w);
        var roiH = (int)(recognition.RoiH * h);
        return Alignment.RoiMatch(w, h, pattern, img, roiX, roiY, roiW, roiH);
    }
}